from torch import nn


class nnModel(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, 192)
        self.gru = nn.GRUCell(192, 192)
        self.output_fc = nn.Linear(192, action_dim)
        self.activation = nn.LeakyReLU(0.05)

    def forward(self, state, hx=None):
        x = self.fc1(state)
        x = self.activation(x)
        hx = self.gru(x, hx)
        action = self.output_fc(self.activation(hx))
        return action, hx
